import asyncio
import json

from loguru import logger

from com.arcfox.middleware import async_redis_middleware as redis
from com.arcfox.util import constant
from com.arcfox.util.date import get_current_timestamp


class RedisTaskManager:
    def __init__(self, task_key):
        self.TASK_KEY = task_key
        self.FAIL_TASK_KEY = f"fail_{self.TASK_KEY}"

    async def add_tasks(self, tasks, unique=True):
        '''
        添加任务
        :param tasks: 任务列表
        :param unique: 是否添加到去重队列
        :return:
        '''
        if not tasks:
            logger.warning("tasks为空!")
            return
        if not isinstance(tasks, list):
            tasks = [tasks]
        if unique:
            await self.__add_unique_tasks(self.TASK_KEY, tasks)
        else:
            await self.__add_tasks(self.TASK_KEY, tasks)

    async def add_fail_tasks(self, tasks, unique=True):
        '''
        任务抓取失败时, 回填任务
        :param tasks: 任务列表
        :param unique: 是否添加到去重队列
        :return:
        '''
        if not isinstance(tasks, list):
            tasks = [tasks]
        if unique:
            await self.__add_unique_tasks(self.FAIL_TASK_KEY, tasks)
        else:
            await self.__add_tasks(self.FAIL_TASK_KEY, tasks)

    @redis.open_redis
    async def add_basic_tasks(self, tasks, client):
        '''
        用于存放基础数据的set, 不设过期时间
        :param key: redis key
        :param tasks:
        :return:
        '''
        if not tasks:
            logger.warning("tasks为空!")
            return
        if not isinstance(tasks, list):
            tasks = [tasks]
        await client.sadd(self.TASK_KEY, *tasks)
        logger.info("task add finish!")

    @redis.open_redis
    async def get_all_basic_tasks(self, format_json, client):
        '''
        获取基础task列表, 只读取, 不删除
        :param key:
        :param format_json:
        :param client:
        :return:
        '''
        tasks = await client.smembers(self.TASK_KEY)
        return self.__format_task(tasks, format_json)

    async def pull_tasks(self, size, unique=True, format_json=True):
        '''
        批量拉取任务
        :param size:
        :param unique: 是否从去重队列拉取(提交任务时也需提交到去重队列)
        :param format_json: 是否将结果转成Json格式
        :return:
        '''
        if unique:
            return await self.__pull_unique_tasks(size, format_json)
        else:
            return await self.__pull_tasks(size, format_json)

    async def task_size(self, unique=True):
        if unique:
            return await self.__unique_task_size()
        else:
            return await self.__task_size()

    @redis.open_redis
    async def __add_tasks(self, task_key, tasks: list, client):
        assert tasks, "task不能为空"
        await client.rpush(task_key, *tasks)
        logger.info("task add finish!")

    @redis.open_redis
    async def __add_unique_tasks(self, task_key, tasks: list, client):
        assert tasks, "tasks不能为空"
        for task in tasks:
            if isinstance(task, dict):
                task = json.dumps(task)
            exist_score = await client.zscore(task_key, task)
            if exist_score:
                continue
            await client.zadd(task_key, constant.MAX_UNIX_TIME - get_current_timestamp(), task)
        logger.info("unique task add finish!")

    @redis.open_redis
    async def __pull_tasks(self, size, format_json, client):
        assert size >= 1, "一次拉取任务不能小于1"
        task_key = self.TASK_KEY
        tasks = await client.lrange(self.TASK_KEY, 0, size - 1)
        if not tasks:
            logger.info("正常任务已拉空, 尝试拉取失败回填任务")
            task_key = self.FAIL_TASK_KEY
            tasks = await client.lrange(self.FAIL_TASK_KEY, 0, size - 1)
        if tasks:
            await client.ltrim(task_key, len(tasks), -1)
        return self.__format_task(tasks, format_json)

    @redis.open_redis
    async def __pull_unique_tasks(self, size, format_json, client):
        assert size >= 1, "一次拉取任务不能小于1"
        task_key = self.TASK_KEY
        task_list = await client.zrevrange(self.TASK_KEY, 0, size - 1)
        if not task_list:
            logger.info("正常任务已拉空, 尝试拉取失败回填任务")
            task_key = self.FAIL_TASK_KEY
            task_list = await client.zrevrange(self.FAIL_TASK_KEY, 0, size - 1)
        if task_list:
            await client.zrem(task_key, *task_list)
        return self.__format_task(task_list, format_json)

    @redis.open_redis
    async def __task_size(self, client):
        normal_task = await client.llen(self.TASK_KEY)
        fail_task = await client.llen(self.FAIL_TASK_KEY)
        return normal_task + fail_task

    @redis.open_redis
    async def __unique_task_size(self, client):
        normal_task = await client.zcard(self.TASK_KEY)
        fail_task = await client.zcard(self.FAIL_TASK_KEY)
        return normal_task + fail_task

    def __format_task(self, tasks, format_json):
        task_list = []
        for task in tasks:
            decode_task = task.decode()
            decode_task = decode_task.replace("\'", "\"")
            if format_json:
                decode_task = self.__try_format_json(decode_task)
            task_list.append(decode_task)
        return task_list

    @staticmethod
    def __try_format_json(task):
        try:
            return json.loads(task)
        except Exception as e:
            logger.exception("任务转json异常", e)
        return task


if __name__ == "__main__":
    manager = RedisTaskManager("basic_test")
    task = {"name": "aa", "age": 10}
    asyncio.run(manager.get_all_basic_tasks(True))
